import os

import albumentations as A
import cv2
import torch
from albumentations.pytorch import ToTensorV2
from timm.models.layers.conv2d_same import conv2d_same

size = 256
transform = A.Compose([
    # A.OneOf([
    #     IsotropicResize(max_side=size, interpolation_down=cv2.INTER_AREA, interpolation_up=cv2.INTER_CUBIC),
    #     IsotropicResize(max_side=size, interpolation_down=cv2.INTER_AREA, interpolation_up=cv2.INTER_LINEAR),
    #     IsotropicResize(max_side=size, interpolation_down=cv2.INTER_LINEAR, interpolation_up=cv2.INTER_LINEAR),
    # ], p=1),
    # A.PadIfNeeded(min_height=size, min_width=size, border_mode=cv2.BORDER_CONSTANT, value=0),
    # A.Normalize(mean=(0., 0., 0.), std=(1., 1., 1.)),
    ToTensorV2(),
])

srm1 = 1 / 4 * torch.tensor([
    [0., 0., 0., 0., 0.],
    [0., -1., 2., -1., 0.],
    [0., 2., -4., 2., 0.],
    [0., -1., 2., -1., 0.],
    [0., 0., 0., 0., 0.]
])

srm2 = 1 / 12 * torch.tensor([
    [-1., 2., -2., 2., -1.],
    [2., -6., 8., -6., 2.],
    [-2., 8., -12., 8., -2.],
    [2., -6., 8., -6., 2.],
    [-1., 2., -2., 2., -1.]
])

srm3 = 1 / 2 * torch.tensor([
    [0., 0., 0., 0., 0.],
    [0., 0., 0., 0., 0.],
    [0., 1., -2., 1., 0.],
    [0., 0., 0., 0., 0.],
    [0., 0., 0., 0., 0.]
])

gaussian = 1 / 9 * torch.tensor([
    [1., 1., 1.],
    [1., 1., 1.],
    [1., 1., 1.],
])

one = 1 / 3 * torch.tensor([
    [0., 0., 0.],
    [0., 1., 0.],
    [0., 0., 0.],
])

sobel_x = torch.tensor([
    [-1., 0., 1.],
    [-2., 0., 2.],
    [-1., 0., 1.]
])

sobel_y = torch.tensor([
    [1., 2., 1.],
    [0., 0., 0.],
    [-1., -2., -1.]
])

scharr_x = torch.tensor([
    [3., 0., -3.],
    [10., 0., -10.],
    [3., 0., -3.]
])

scharr_y = torch.tensor([
    [3., 10., 3.],
    [0., 0., 0.],
    [-3., -10., -3.]
])


def srm_filter(x):
    srm_weight = torch.stack((srm1, srm2, srm3))
    # srm_weight = torch.stack((srm2, srm2, srm2))
    # srm_weight = torch.stack((gaussian, gaussian, gaussian))
    # srm_weight = torch.stack((one, one, one))
    srm_weight = srm_weight.unsqueeze(1)
    srm_weight = srm_weight.repeat(3, 1, 1, 1)
    x = conv2d_same(x, srm_weight, stride=(1, 1), groups=1)
    return x


def gaussian_filter(x):
    B, C, H, W = x.shape
    gaussian_weight = gaussian.view(1, 1, 3, 3).repeat(3, 1, 1, 1)
    x = conv2d_same(x, gaussian_weight, stride=(1, 1), groups=C)
    return x


def image_gradient_enhance(image):
    transformed_image = transform(image=image)['image'].to(torch.float32)
    transformed_image = torch.unsqueeze(transformed_image, dim=0)

    # transformed_image = srm_filter(transformed_image)
    weight = sobel_x
    weight = weight.repeat(3, 1, 1, 1)
    gradient_x = conv2d_same(transformed_image, weight, stride=(2, 2), groups=3)

    weight = sobel_y
    weight = weight.repeat(3, 1, 1, 1)
    gradient_y = conv2d_same(transformed_image, weight, stride=(2, 2), groups=3)

    sobel_image = torch.sqrt(torch.square(gradient_x) + torch.square(gradient_y))

    # print(sobel_image.shape)
    sobel_image = sobel_image.squeeze().permute(1, 2, 0).numpy()
    return sobel_image


def main():
    """
    E:\\Datasets\\test\\ff++_nt\\402__10_0.png
    E:\\Datasets\\test\\ff++_nt\\402_453__10_0.png
    """
    # image_path = 'E:\\workspace-pycharm\\deepfakes\\sifdnet-main\\data\\img\\Deepfakes0_0.png'
    root_path = 'E:\\workspace-pycharm\\deepfakes\\sifdnet-main\\plot\\img'
    image_folder = os.path.join(root_path, 'image-gradient')
    for image_name in os.listdir(image_folder):
        image_path = os.path.join(image_folder, image_name)
        image = cv2.imread(image_path, cv2.IMREAD_COLOR)
        image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)

        # transformed_image = (transformed_image - np.min(transformed_image)) / \
        #                     (np.max(transformed_image) - np.min(transformed_image))
        # transformed_image = transformed_image.squeeze().numpy()
        # transformed_image = cv2.cvtColor(transformed_image, cv2.COLOR_RGB2BGR)
        sobel_image = image_gradient_enhance(image)
        target_file = os.path.join(image_folder, '{}_{}.png'.format(image_name[:-4], 'sobel'))
        cv2.imwrite(target_file, sobel_image)


if __name__ == '__main__':
    main()
